from ultralytics import YOLO
import cv2

model = YOLO("runs/detect/train-yolo8l-150t-alldata/weights/best.pt")


def change_model(model_path):
    global model
    model = YOLO(model_path)


def detect_beetles(image_path):
    results = model(image_path)
    img = cv2.imread(image_path)
    boxes = results[0].boxes

    beetle_count = len(boxes)
    bbox_list = []
    confidences = []

    for box in boxes:
        xyxy = box.xyxy[0].cpu().numpy().astype(int)
        conf = box.conf.item()
        confidences.append(conf)
        bbox_list.append((xyxy, conf))
        x1, y1, x2, y2 = xyxy
        cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
        cv2.putText(
            img,
            f"{conf:.2f}",
            (x1, y1 - 10),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.5,
            (0, 255, 0),
            2,
        )

    avg_conf = float(sum(confidences)) / beetle_count if beetle_count else 0
    return img, beetle_count, avg_conf, bbox_list
